# author zybz 2022.3.16
# step3 生成配套的 yaml 文件
import yaml
import os
import shutil


class get_yaml():
    """generate yaml file
    -- if you want to copy the yaml file to 'yolov5/data', you need to set copy=True and your copy_dst_path like the example above
    """
    def __init__(self, classes, copy=False, copy_dst_path='robocon-vision/data'):
        """make yaml init
            if you want to copy the yaml file to 'yolov5/data', you need to set copy=True and your copy_dst_path like the example above

        Args:
            classes (list): _Dataset labels_
            copy (bool, optional): Copy yaml file to yolov5. Defaults to False.
            copy_dst_path (str, optional): copy destination path. Defaults to 'robocon-vision/data'.
        """
        if classes is None:
            print("please input classes")
            return
        self.classes = classes
        self.datasets_name = os.getcwd().split('\\')[-2]  # 此数据集文件夹名字
        self.copy = copy
        self.copy_dst_path = copy_dst_path
        self.run()

    def get_data(self):
        """Yaml file content formatting and filling

        Returns:
            _type_: _description_
        """
        data = {}
        relative_path = '../datasets/' + self.datasets_name  # 该路径用于train.py中

        data['names'] = list(self.classes)
        data['nc'] = len(self.classes)
        data['test'] = 'test.txt'
        data['val'] = 'val.txt'
        data['train'] = 'train.txt'
        data['path'] = relative_path
        return data

    def run(self):
        """Run all functions of this class
        """
        file_path = '../' + self.datasets_name + '.yaml'
        with open(file_path, 'w', encoding='utf-8') as f:
            yaml.dump_all(documents=[self.get_data()], stream=f, allow_unicode=True)
        # add comments to recommend to use
        with open(file_path, 'a', encoding='utf-8') as f:
            f.write("# Usage: \n")
            f.write(
                "#    $ nohup python train.py --data newblock_0316.yaml --weights yolov5s.pt --img 640 --device 0 --epochs 600 --batch-size 64 --patience 600 > logs/log_newblock_0316.txt 2>&1 &"
            )
        # copy *.yaml to yolov5 folder 
        if self.copy is True and self.copy_dst_path is not None:
            self.copy_yaml_To_yolov5_folder(destination_path=self.copy_dst_path)

        print("Step3 make yaml successfully")

    def copy_yaml_To_yolov5_folder(self, destination_path='robocon-vision/data'):
        """Copy yaml file to yolov5

        Args:
            destination_path (str, optional): copy destination path. Defaults to 'robocon-vision/data'.
        """
        src_path = '../'
        src_file = src_path + self.datasets_name + '.yaml'
        dst_path = '../../../' + destination_path + '/'  # 进入与 datasets 同级的 yolov5 文件夹中
        dst_file = dst_path + self.datasets_name + '.yaml'
        shutil.copyfile(src_file, dst_file)


if __name__ == "__main__":
    classes = ["fire", "fall"]
    demo = get_yaml(classes, copy=True, copy_dst_path='robocon-vision/data')
